/*
 * Copyright (c) 2011-2013 Jarek Sacha. All Rights Reserved.
 *
 * Author's e-mail: jpsacha at gmail.com
 */
package br.com.manchini.stereocv.utils;

import com.googlecode.javacv.cpp.opencv_core;
import com.googlecode.javacv.cpp.opencv_features2d;
import com.googlecode.javacv.cpp.opencv_nonfree;
import java.util.ArrayList;
import static com.googlecode.javacv.cpp.opencv_calib3d.*;
import static com.googlecode.javacv.cpp.opencv_features2d.*;
import static com.googlecode.javacv.cpp.opencv_core.*;

public class RobustMatcher {

    private opencv_nonfree.SURF detector;
    
    private opencv_features2d.DescriptorExtractor extractor;
    private float ratio;
    private boolean refineF;
    private double minDistanceToEpipolar;
    private double confidenceLevel;

    /**
     * Robust matcher used by examples Ex3ComputeFundamentalMatrix and
     * Ex5Homography.
     *
     * See Chapter 9 page 233.
     *
     * @param detector Feature point detector
     * @param extractor Feature descriptor extractor
     * @param ratio Max ratio between 1st and 2nd NN
     * @param refineF If `true` will refine the F matrix
     * @param minDistanceToEpipolar Min distance to epipolar
     * @param confidenceLevel Confidence level (probability)
     */
    public void RobustMatcher() {

        detector = new opencv_nonfree.SURF(100);
        extractor = opencv_features2d.DescriptorExtractor.create("SURF");
        ratio = 0.65f;
        refineF = true;
        minDistanceToEpipolar = 3.0;
        confidenceLevel = 0.99;
        

    }

    public RobustMatcher(double confidenceLevel, double minDistanceToEpipolar, float ratio, opencv_nonfree.SURF detector) {
        this.detector = detector;
        this.ratio = ratio;
        this.minDistanceToEpipolar = minDistanceToEpipolar;
        this.confidenceLevel = confidenceLevel;
        extractor = opencv_features2d.DescriptorExtractor.create("SURF");
        
        
    }

    
    /**
     * Match feature points using symmetry test and RANSAC
     *
     * @return fundamental matrix.
     */
    public Result matchImages(opencv_core.IplImage image1, opencv_core.IplImage image2, IplImage mascaraEsq, IplImage mascaraDir) {

        // 1a. Detection of the SURF features
        KeyPoint keyPoints1 = new KeyPoint();
        KeyPoint keyPoints2 = new KeyPoint();
        detector.detect(image1, mascaraEsq, keyPoints1);
        detector.detect(image2, mascaraDir, keyPoints2);
        System.out.println("Number of SURF points (1): " + keyPoints1.capacity());
        System.out.println("Number of SURF points (2): " + keyPoints2.capacity()); // 1b. Extraction of the SURF descriptors
        CvMat descriptors1 = new CvMat(null);
        CvMat descriptors2 = new CvMat(null);
        extractor.compute(image1, keyPoints1, descriptors1);
        extractor.compute(image2, keyPoints2, descriptors2);
        
        System.out.println("descriptor matrix size: " + descriptors1.rows() + " by " + descriptors1.cols()); // 2. Match the two image descriptors
        // Construction of the matcher
        opencv_features2d.BFMatcher matcher = new opencv_features2d.BFMatcher(opencv_core.NORM_L2, false); // from image 1 to image 2
        // based on k nearest neighbours (with k=2)
        opencv_features2d.DMatchVectorVector matches1 = new DMatchVectorVector();
        matcher.knnMatch(
                descriptors1, descriptors2,
                matches1, // vector of matches (up to 2 per entry)
                2, // return 2 nearest neighbours
                null, // mask
                false // compact result, used when mask != null
                ); // from image 2 to image 1
        // based on k nearest neighbours (with k=2)
        opencv_features2d.DMatchVectorVector matches2 = new DMatchVectorVector();
        matcher.knnMatch(
                descriptors2, descriptors1,
                matches2, // vector of matches (up to 2 per entry)
                2, // return 2 nearest neighbours
                null, // mask
                false // compact result, used when mask != null
                );

        System.out.println("Number of matched points 1->2: " + matches1.size());
        System.out.println("Number of matched points 2->1: " + matches2.size()); // 3. Remove matches for which NN ratio is > than threshold
        //  ... clean image 1 -> image 2 matches
        ArrayList<DMatch[]> matches1WithRatio = ratioTest(matches1);
        System.out.println("Number of matched points 1->2 with ratio: " + matches1WithRatio.size()); //  ... clean image 2 -> image 1 matches
        ArrayList<DMatch[]> matches2WithRatio = ratioTest(matches2);
        System.out.println("Number of matched points 2->1 with ratio: " + matches2WithRatio.size()); // 4. Remove non-symmetrical matches
        ArrayList<DMatch> symMatches = symmetryTest(matches1WithRatio, matches2WithRatio);
        System.out.println("Number of matched points (symmetry test): " + symMatches.size()); // 5. Validate matches using RANSAC
//        val(matches, fundamentalMatrix) = ransacTest(symMatches, keyPoints1, keyPoints2);
        Object[] retorno = ransacTest(symMatches, keyPoints1, keyPoints2);
        ArrayList<DMatch> matches = (ArrayList<DMatch>) retorno[0];
        CvMat fundamentalMatrix = (CvMat) retorno[1];

        return new Result(matches, keyPoints1, keyPoints2, fundamentalMatrix);
    }

    /**
     * Filter matches for which NN ratio is > than threshold, also remove
     * non-matches, if present in the input.
     *
     * @param matches collection of matches that will be filtered.
     * @return the number of removed points (corresponding entries being
     * cleared, i.e. size will be 0)
     */
    private ArrayList<DMatch[]> ratioTest(DMatchVectorVector matches) {

        // Find matches that need to be removed
        ArrayList<DMatch[]> destArray = new ArrayList();

        for (int i = 0; i < matches.size(); i++) {
            // if 2 NN has been identified
            if (matches.size(i) > 1) {
                if (matches.get(i, 0).distance() / matches.get(i, 1).distance() <= ratio) {
                    destArray.add(new DMatch[]{matches.get(i, 0), matches.get(i, 1)});
                }
            }
        }

        return destArray;
    }

    /**
     * Insert symmetrical matches in returned array.
     */
    private ArrayList<DMatch> symmetryTest(ArrayList<DMatch[]> matches1,
            ArrayList<DMatch[]> matches2) {

        ArrayList<DMatch> destSeq = new ArrayList<DMatch>();

        // For all matches image 1 -> image 2
        for (DMatch[] m1 : matches1) {
            DMatch m11 = m1[0];
            boolean para = false;
            for (DMatch[] m2 : matches2) {
                DMatch m21 = m2[0];
                if (m11.queryIdx() == m21.trainIdx() && m21.queryIdx() == m11.trainIdx()) {
                    destSeq.add(new DMatch(m11.queryIdx(), m11.trainIdx(), m11.distance()));
                    para = true;
                }
            }
        }

        return destSeq;
    }

 

    /**
     * Identify good matches using RANSAC
     *
     * @param matches input matches
     * @return surviving matches and the fundamental matrix
     */
    public Object[] ransacTest(ArrayList<DMatch> matches, KeyPoint keyPoints1, KeyPoint keyPoints2) {

        // Convert keypoints into Point2f
        opencv_core.CvPoint2D32f[] points = MatcherUtils.toCvPoint2D32f(matches, keyPoints1, keyPoints2); // Compute F matrix using RANSAC
        opencv_core.CvPoint2D32f point1 = points[0];
        opencv_core.CvPoint2D32f point2 = points[1];
        CvMat fundamentalMatrix = CvMat.create(3, 3, opencv_core.CV_32F);
        CvMat pointStatus = CvMat.create(matches.size(), 1, opencv_core.CV_8U);
        cvFindFundamentalMat(
                MatcherUtils.toCvMat(point1) /*  points in first image */,
                MatcherUtils.toCvMat(point2) /*  points in second image */,
                fundamentalMatrix /* output */,
                CV_FM_RANSAC /* RANSAC method */,
                minDistanceToEpipolar, /* distance to epipolar plane */
                confidenceLevel /* confidence probability */,
                pointStatus /* match status (inlier ou outlier) */); // extract the surviving (inliers) matches
        ArrayList<DMatch> outMatches = new ArrayList<DMatch>();

        for (int i = 0; i < pointStatus.rows(); i++) {
            Long inlier = Math.round(pointStatus.get(i));
            if (inlier != 0) {
                outMatches.add(matches.get(i));
            }
        }

        System.out.println("Number of matched points (after cleaning): " + outMatches.size());

        if (refineF) {
            // The F matrix will be recomputed with all accepted matches
            points = MatcherUtils.toCvPoint2D32f(outMatches, keyPoints1, keyPoints2);
            point1 = points[0];
            point2 = points[1];

            // Compute 8-point F from all accepted matches
            cvFindFundamentalMat(
                    MatcherUtils.toCvMat(point1) /* points in first image */,
                    MatcherUtils.toCvMat(point2) /* points in second image */,
                    fundamentalMatrix /* output */,
                    CV_FM_8POINT /* 8-point method */,
                    minDistanceToEpipolar, /* distance to epipolar plane, only used with CV_FM_RANSAC */
                    confidenceLevel /* confidence probability, only used with CV_FM_RANSAC */,
                    null);
        }

        return new Object[]{outMatches, fundamentalMatrix};


    }

    public class Result {

        ArrayList<DMatch> matches;
        KeyPoint keyPoints1;
        KeyPoint keyPoints2;
        CvMat fundamentalMatrix;

        public Result(ArrayList<DMatch> matches, KeyPoint keyPoints1, KeyPoint keyPoints2, CvMat fundamentalMatrix) {
            this.matches = matches;
            this.keyPoints1 = keyPoints1;
            this.keyPoints2 = keyPoints2;
            this.fundamentalMatrix = fundamentalMatrix;
        }

        public ArrayList<DMatch> getMatches() {
            return matches;
        }

        public KeyPoint getKeyPoints1() {
            return keyPoints1;
        }

        public KeyPoint getKeyPoints2() {
            return keyPoints2;
        }

        public CvMat getFundamentalMatrix() {
            return fundamentalMatrix;
        }
        
    }
}